{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "72b999e0-1dd0-4586-9577-91046598468c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: pandas in ./.conda/envs/thinp/lib/python3.11/site-packages (2.2.2)\n",
      "Requirement already satisfied: numpy>=1.23.2 in ./.conda/envs/thinp/lib/python3.11/site-packages (from pandas) (1.26.4)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in ./.conda/envs/thinp/lib/python3.11/site-packages (from pandas) (2.9.0.post0)\n",
      "Requirement already satisfied: pytz>=2020.1 in ./.conda/envs/thinp/lib/python3.11/site-packages (from pandas) (2024.1)\n",
      "Requirement already satisfied: tzdata>=2022.7 in ./.conda/envs/thinp/lib/python3.11/site-packages (from pandas) (2024.1)\n",
      "Requirement already satisfied: six>=1.5 in ./.conda/envs/thinp/lib/python3.11/site-packages (from python-dateutil>=2.8.2->pandas) (1.16.0)\n"
     ]
    }
   ],
   "source": [
    "!pip install pandas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d1dc0498-4efd-4bad-b9aa-781650dff064",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: spacy in ./.conda/envs/thinp/lib/python3.11/site-packages (3.7.5)\n",
      "Requirement already satisfied: spacy-legacy<3.1.0,>=3.0.11 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (3.0.12)\n",
      "Requirement already satisfied: spacy-loggers<2.0.0,>=1.0.0 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (1.0.5)\n",
      "Requirement already satisfied: murmurhash<1.1.0,>=0.28.0 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (1.0.10)\n",
      "Requirement already satisfied: cymem<2.1.0,>=2.0.2 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (2.0.8)\n",
      "Requirement already satisfied: preshed<3.1.0,>=3.0.2 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (3.0.9)\n",
      "Requirement already satisfied: thinc<8.3.0,>=8.2.2 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (8.2.4)\n",
      "Requirement already satisfied: wasabi<1.2.0,>=0.9.1 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (1.1.3)\n",
      "Requirement already satisfied: srsly<3.0.0,>=2.4.3 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (2.4.8)\n",
      "Requirement already satisfied: catalogue<2.1.0,>=2.0.6 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (2.0.10)\n",
      "Requirement already satisfied: weasel<0.5.0,>=0.1.0 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (0.4.1)\n",
      "Requirement already satisfied: typer<1.0.0,>=0.3.0 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (0.12.3)\n",
      "Requirement already satisfied: tqdm<5.0.0,>=4.38.0 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (4.66.4)\n",
      "Requirement already satisfied: requests<3.0.0,>=2.13.0 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (2.32.2)\n",
      "Requirement already satisfied: pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (2.7.3)\n",
      "Requirement already satisfied: jinja2 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (3.1.4)\n",
      "Requirement already satisfied: setuptools in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (69.5.1)\n",
      "Requirement already satisfied: packaging>=20.0 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (23.2)\n",
      "Requirement already satisfied: langcodes<4.0.0,>=3.2.0 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (3.4.0)\n",
      "Requirement already satisfied: numpy>=1.19.0 in ./.conda/envs/thinp/lib/python3.11/site-packages (from spacy) (1.26.4)\n",
      "Requirement already satisfied: language-data>=1.2 in ./.conda/envs/thinp/lib/python3.11/site-packages (from langcodes<4.0.0,>=3.2.0->spacy) (1.2.0)\n",
      "Requirement already satisfied: annotated-types>=0.4.0 in ./.conda/envs/thinp/lib/python3.11/site-packages (from pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4->spacy) (0.7.0)\n",
      "Requirement already satisfied: pydantic-core==2.18.4 in ./.conda/envs/thinp/lib/python3.11/site-packages (from pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4->spacy) (2.18.4)\n",
      "Requirement already satisfied: typing-extensions>=4.6.1 in ./.conda/envs/thinp/lib/python3.11/site-packages (from pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4->spacy) (4.11.0)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in ./.conda/envs/thinp/lib/python3.11/site-packages (from requests<3.0.0,>=2.13.0->spacy) (2.0.4)\n",
      "Requirement already satisfied: idna<4,>=2.5 in ./.conda/envs/thinp/lib/python3.11/site-packages (from requests<3.0.0,>=2.13.0->spacy) (2.10)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in ./.conda/envs/thinp/lib/python3.11/site-packages (from requests<3.0.0,>=2.13.0->spacy) (2.2.1)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in ./.conda/envs/thinp/lib/python3.11/site-packages (from requests<3.0.0,>=2.13.0->spacy) (2024.6.2)\n",
      "Requirement already satisfied: blis<0.8.0,>=0.7.8 in ./.conda/envs/thinp/lib/python3.11/site-packages (from thinc<8.3.0,>=8.2.2->spacy) (0.7.11)\n",
      "Requirement already satisfied: confection<1.0.0,>=0.0.1 in ./.conda/envs/thinp/lib/python3.11/site-packages (from thinc<8.3.0,>=8.2.2->spacy) (0.1.5)\n",
      "Requirement already satisfied: click>=8.0.0 in ./.conda/envs/thinp/lib/python3.11/site-packages (from typer<1.0.0,>=0.3.0->spacy) (8.1.7)\n",
      "Requirement already satisfied: shellingham>=1.3.0 in ./.conda/envs/thinp/lib/python3.11/site-packages (from typer<1.0.0,>=0.3.0->spacy) (1.5.4)\n",
      "Requirement already satisfied: rich>=10.11.0 in ./.conda/envs/thinp/lib/python3.11/site-packages (from typer<1.0.0,>=0.3.0->spacy) (13.7.1)\n",
      "Requirement already satisfied: cloudpathlib<1.0.0,>=0.7.0 in ./.conda/envs/thinp/lib/python3.11/site-packages (from weasel<0.5.0,>=0.1.0->spacy) (0.18.1)\n",
      "Requirement already satisfied: smart-open<8.0.0,>=5.2.1 in ./.conda/envs/thinp/lib/python3.11/site-packages (from weasel<0.5.0,>=0.1.0->spacy) (7.0.4)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in ./.conda/envs/thinp/lib/python3.11/site-packages (from jinja2->spacy) (2.1.3)\n",
      "Requirement already satisfied: marisa-trie>=0.7.7 in ./.conda/envs/thinp/lib/python3.11/site-packages (from language-data>=1.2->langcodes<4.0.0,>=3.2.0->spacy) (1.2.0)\n",
      "Requirement already satisfied: markdown-it-py>=2.2.0 in ./.conda/envs/thinp/lib/python3.11/site-packages (from rich>=10.11.0->typer<1.0.0,>=0.3.0->spacy) (3.0.0)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in ./.conda/envs/thinp/lib/python3.11/site-packages (from rich>=10.11.0->typer<1.0.0,>=0.3.0->spacy) (2.15.1)\n",
      "Requirement already satisfied: wrapt in ./.conda/envs/thinp/lib/python3.11/site-packages (from smart-open<8.0.0,>=5.2.1->weasel<0.5.0,>=0.1.0->spacy) (1.16.0)\n",
      "Requirement already satisfied: mdurl~=0.1 in ./.conda/envs/thinp/lib/python3.11/site-packages (from markdown-it-py>=2.2.0->rich>=10.11.0->typer<1.0.0,>=0.3.0->spacy) (0.1.2)\n"
     ]
    }
   ],
   "source": [
    "!pip install spacy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4132ef9b-bab7-4dfa-be33-a07d23879741",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: requests in ./.conda/envs/thinp/lib/python3.11/site-packages (2.32.2)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in ./.conda/envs/thinp/lib/python3.11/site-packages (from requests) (2.0.4)\n",
      "Requirement already satisfied: idna<4,>=2.5 in ./.conda/envs/thinp/lib/python3.11/site-packages (from requests) (2.10)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in ./.conda/envs/thinp/lib/python3.11/site-packages (from requests) (2.2.1)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in ./.conda/envs/thinp/lib/python3.11/site-packages (from requests) (2024.6.2)\n"
     ]
    }
   ],
   "source": [
    "!pip install requests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "61037eca-d132-4723-b8c1-5cd4d9b94ec3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import spacy\n",
    "\n",
    "df = pd.read_excel(\"answer.xlsx\")\n",
    "\n",
    "result_df = df[['Value']]\n",
    "\n",
    "result_list = result_df['Value'].tolist()\n",
    "\n",
    "result = []\n",
    "for i in range(len(result_list)):\n",
    "    result.append(str(result_list[i]).split(\", \"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "01ddec8b-4c78-46a9-a45a-1f7c4adc7a77",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "with open('test.jsonl', 'r') as f:\n",
    "    data = [json.loads(line) for line in f]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5f7e82b4-057b-4d97-94d3-2e4524813db4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mapping(ans, relation):\n",
    "    mapping_list = []\n",
    "    for i in range(len(ans)):\n",
    "        if ans[i] == \"nan\":\n",
    "            break\n",
    "        elif relation == \"seriesHasNumberOfEpisodes\":\n",
    "            mapping_list.append(ans[i])\n",
    "        else:\n",
    "            answer = ans[i]\n",
    "            if \": \" in answer:\n",
    "                answer_split = answer.split(\": \")\n",
    "                for item in answer_split:\n",
    "\n",
    "                    entity_wikidata = search_entity_wikidata(item)\n",
    "                        \n",
    "                    if entity_wikidata != None:\n",
    "                        instance = get_instance_of(entity_wikidata['QID'])\n",
    "                        entity_wikidata['Instance of'] = instance\n",
    "                        \n",
    "                        if entity_wikidata['QID'] not in mapping_list:\n",
    "                            if relation == \"awardWonBy\":\n",
    "                                if \"human\" in entity_wikidata['Instance of']:\n",
    "                                    mapping_list.append(entity_wikidata['QID'])\n",
    "                            elif relation == \"companyTradesAtStockExchange\":\n",
    "                                if \"stock exchange\" in entity_wikidata['Instance of']:\n",
    "                                    mapping_list.append(entity_wikidata['QID'])\n",
    "                            elif relation == \"countryLandBordersCountry\":\n",
    "                                if \"country\" in entity_wikidata['Instance of']:\n",
    "                                    mapping_list.append(entity_wikidata['QID'])\n",
    "                            else: # personHasCityOfDeath\n",
    "                                for i in range(len(entity_wikidata['Instance of'])):\n",
    "                                    if \"city\" in entity_wikidata['Instance of'][i]:\n",
    "                                        mapping_list.append(entity_wikidata['QID'])\n",
    "                                        break\n",
    "            else:\n",
    "  \n",
    "                entity_wikidata = search_entity_wikidata(answer)\n",
    "    \n",
    "                if entity_wikidata != None:\n",
    "                    instance = get_instance_of(entity_wikidata['QID'])\n",
    "                    entity_wikidata['Instance of'] = instance\n",
    "                    if entity_wikidata['QID'] not in mapping_list:\n",
    "                        if relation == \"awardWonBy\":\n",
    "                            if \"human\" in entity_wikidata['Instance of']:\n",
    "                                mapping_list.append(entity_wikidata['QID'])\n",
    "                        elif relation == \"companyTradesAtStockExchange\":\n",
    "                            if \"stock exchange\" in entity_wikidata['Instance of']:\n",
    "                                mapping_list.append(entity_wikidata['QID'])\n",
    "                        elif relation == \"countryLandBordersCountry\":\n",
    "                            if \"country\" in entity_wikidata['Instance of']:\n",
    "                                mapping_list.append(entity_wikidata['QID'])\n",
    "                        else: # personHasCityOfDeath\n",
    "                            for i in range(len(entity_wikidata['Instance of'])):\n",
    "                                if \"city\" in entity_wikidata['Instance of'][i]:\n",
    "                                    mapping_list.append(entity_wikidata['QID'])\n",
    "                                    break\n",
    "\n",
    "    return mapping_list\n",
    "                    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "bfeef877-2e12-441b-a7bf-d478ddf1a4a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "import time\n",
    "\n",
    "def search_entity_wikidata(entity_name, language='en'):\n",
    "    url = 'https://www.wikidata.org/w/api.php'\n",
    "    params = {\n",
    "        'action': 'wbsearchentities',\n",
    "        'search': entity_name,\n",
    "        'language': language,\n",
    "        'format': 'json'\n",
    "    }\n",
    "    \n",
    "    response = requests.get(url, params=params)\n",
    "        \n",
    "    data = response.json()\n",
    "    # Check if there are any search results\n",
    "    if data['search']:\n",
    "        # Return only QID and Entity name\n",
    "        results = [{'QID': result['id'], 'EntityName': result['label']} for result in data['search']]\n",
    "        # Always choose the 1st one\n",
    "        return results[0]\n",
    "        \n",
    "    else: return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6d929228-fb55-4b50-8134-12e93d181f3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_instance_of(qid):\n",
    "    query = f\"\"\"\n",
    "    SELECT ?instanceOfLabel WHERE {{\n",
    "      wd:{qid} wdt:P31 ?instanceOf .\n",
    "      SERVICE wikibase:label {{ bd:serviceParam wikibase:language \"[AUTO_LANGUAGE],en\". }}\n",
    "    }}\n",
    "    \"\"\"\n",
    "    \n",
    "    url = 'https://query.wikidata.org/sparql'\n",
    "    headers = {\n",
    "        'Accept': 'application/sparql-results+json',\n",
    "    }\n",
    "    \n",
    "    response = requests.get(url, params={'query': query}, headers=headers)\n",
    "    \n",
    "    if response.status_code == 200:\n",
    "        results = response.json()['results']['bindings']\n",
    "        return [result['instanceOfLabel']['value'] for result in results]\n",
    "    elif response.status_code == 429:\n",
    "        time.sleep(10)  \n",
    "        return get_instance_of(qid)\n",
    "    else:\n",
    "        response.raise_for_status()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e9a9c451-051b-4598-838e-3c6128a9b223",
   "metadata": {},
   "outputs": [],
   "source": [
    "Result_mapping = []\n",
    "\n",
    "for i in range(len(data)):\n",
    "    Result_mapping.append(mapping(result[i], data[i]['Relation']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "691d4c33-35a0-41b5-9910-7c102c75703d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Saving Result\n",
    "\n",
    "df = pd.DataFrame({'Result': [str(lst) for lst in Result_mapping]})\n",
    "\n",
    "df.to_excel('mapping.xlsx', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "7b624cde-7511-4023-915c-dafa8392fa8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = []\n",
    "for i in range(len(data)):\n",
    "    list = {}\n",
    "    list[\"SubjectEntity\"] = data[i][\"SubjectEntity\"]\n",
    "    list[\"Relation\"] = data[i][\"Relation\"]\n",
    "    list[\"ObjectEntitiesID\"] = Result_mapping[i]\n",
    "    predictions.append(list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "da1c513b-857f-42c1-9519-42b3cdeb2561",
   "metadata": {},
   "outputs": [],
   "source": [
    "fp = \"predictions.jsonl\"\n",
    "\n",
    "with open(fp, \"w\") as f:\n",
    "    for pred in predictions:\n",
    "        f.write(json.dumps(pred) + \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39f2788c-4de6-486f-aab6-a22785d0c11a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
